# -*- coding: utf-8 -*-
# @Author: Alex Gang
# @Date: 2025/6/5 09:07
# @Description:

import torch.optim as optim
import torch
from torch import nn
import torch.nn.functional as F
import math

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, d_model):
        #vocabsize是词汇表大小，d_model是嵌入维度
        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)

class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_len, device='cuda'):
        super(PositionalEmbedding, self).__init__()
        # maxlen理解为句子的最大长度
        self.encoding = torch.zeros(max_len, d_model, device=device)
        # encoding过程不要求梯度更新
        self.encoding.requires_grad = False
        # pos为[0,1,2,,,,max_len]
        pos = torch.arange(0, max_len, device=device)
        # pos 维度从一维变成 （maxlen,1）
        pos = pos.float().unsqueeze(dim=1)
        # _2i is tensor [0,2,4,6,8 ....], 一共有 d_model/2 个元素
        _2i = torch.arange(0, d_model,step=2, device=device).float()
        # 显然位置编码值是0-1之间，也就是小数部分
        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))

    def forward(self, x):
        # 位置编码返回的是一个tensor，维度为[seq_len, d_model]
        batch_size, seq_len = x.size()
        return self.encoding[:seq_len, :]

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, drop_prob, device='cuda'):
        super(TransformerEmbedding, self).__init__()
        self.token_embedding = TokenEmbedding(vocab_size, d_model)
        self.positional_embedding = PositionalEmbedding(d_model, max_len, device)
        self.drop_out = nn.Dropout(drop_prob)

    def forward(self, x):
        # x: [batch_size, seq_len]
        token_emb = self.token_embedding(x)  # [batch_size, seq_len, d_model]
        pos_emb = self.positional_embedding(x)  # [seq_len, d_model]
        return self.drop_out(token_emb + pos_emb) # [batch_size, seq_len, d_model]


class MutiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head, drop_prob=0.1):
        super(MutiHeadAttention, self).__init__()
        # 模型维度 和 头数
        self.n_head = n_head
        self.d_model = d_model
        # q k v
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_combine = nn.Linear(d_model, d_model)
        self.softmax = nn.Softmax(dim=-1)


    def forward(self, query, key, value, mask=None):
        batch, time ,dimension = query.size()
        n_d = self.d_model // self.n_head
        q,k,v = self.w_q(query), self.w_k(key), self.w_v(value)
        q = q.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)  # [batch, n_head, time, n_d]
        k = k.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)
        v = v.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3)
        # scores = query@k.transpose(2,3) / math.sqrt(n_d)  # [batch, n_head, time, time]
        scores = query @ k.transpose(-2, -1) / math.sqrt(n_d)  # [batch, n_head, time, time]
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -10000)
        scores = nn.Softmax(scores)@value
        scores = scores.permute(0, 2, 1, 3).contiguous().view(batch, time, dimension)
        out = self.w_combine(scores)
        return out

class LayerNorm(nn.Module):
    # eps初始稳定性
    def __init__(self, d_model, eps=1e-12):
        super(LayerNorm, self).__init__()
        # nn.Parameter作用是注册一个可学习的参数，也就是说gamma和beta会在训练过程中被更新
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        # 取x最后一个维度的均值和方差
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, unbiased=False, keepdim=True)
        out = (x - mean) / (torch.sqrt(var + self.eps))
        return self.gamma * out + self.beta

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, hidden , dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, hidden)
        self.fc2 = nn.Linear(hidden, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x=self.dropout (x)
        x = self.fc2(x)
        return x

class EncoderLayer(nn.Module):
    def __init__(self, d_model, ffn_hidden, n_head,  dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.attention = MutiHeadAttention(d_model, n_head)
        self.norm1 = LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.ffn = PositionwiseFeedForward(d_model, ffn_hidden, dropout)

    def forward(self, x, mask=None):
        _x = x
        x = self.attention(x, x, x, mask)
        x = self.dropout1(x)
        x = self.norm1(x + _x)
        _x = x
        x = self.ffn(x)
        x = self.dropout2(x)
        x = self.norm1(x + _x)
        return x

class Encoder(nn.Module):
    def __init__(self, enc_voc_size, max_len, d_model, ffn_hidden ,n_head,n_layer,device, dropout=0.1):
        super(Encoder, self).__init__()
        self.embedding = TransformerEmbedding(enc_voc_size, d_model, max_len,dropout)
        self.layers = nn.ModuleList(
            [
                EncoderLayer(d_model, ffn_hidden, n_head, dropout)
                for _ in range(n_layer)
            ]
        )

    def forward(self, x, s_mask):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x, s_mask)
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, ffn_hidden, n_head, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.attention = MutiHeadAttention(d_model, n_head)
        self.norm1 = LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.cross_attention = MutiHeadAttention(d_model, n_head)
        self.norm2 = LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.ffn = PositionwiseFeedForward(d_model, d_model, dropout)
        self.norm3 = LayerNorm(d_model)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, dec,enc,t_mask, s_mask):
        _x = dec
        x = self.attention(dec, dec, dec, t_mask)
        x = self.dropout1(x)
        x = self.norm1(x + _x)
        _x = x
        x = self.cross_attention(x, enc, enc, s_mask)
        x = self.dropout2(x)
        x = self.norm2(x + _x)
        x = self.ffn(x)
        x = self.dropout3(x)
        x = self.norm3(x + _x)
        return x

class Decoder(nn.Module):
    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden,n_head,n_layer,device, dropout=0.1):
        super(Decoder, self).__init__()
        self.embedding = TransformerEmbedding(dec_voc_size,d_model, max_len, dropout,device)
        self.layers = nn.ModuleList(
            [
                DecoderLayer(d_model, ffn_hidden, n_head, dropout)
                for _ in range(n_layer)
            ]
        )
        self.fc = nn.Linear(d_model, dec_voc_size)

    def forward(self, dec, enc, t_mask, s_mask):
        dec = self.embedding(dec)
        for layer in self.layers:
            x = layer(x, s_mask)
        dec = self.fc(dec)
        return dec


class Transformer(nn.Module):
    def __init__(self, src_pad_idx, trg_pad_idx, enc_voc_size, dec_voc_size, d_model,max_len, n_heads, ffn_hidden,n_layer,device, dropout=0.1):
        super(Transformer, self).__init__()
        # transformer整体依然是 encoder-decoder架构
        self.encoder = Encoder(enc_voc_size,max_len,d_model,ffn_hidden ,n_heads,n_layer,device,dropout)
        self.decoder = Decoder(dec_voc_size, max_len, d_model, ffn_hidden,n_heads,n_layer,device,dropout)
        self.src_pad_ix = src_pad_idx
        self.trg_pad_ix = trg_pad_idx
        self.device = device
        self.max_len = max_len
        self.n_heads = n_heads

    def make_pad_mask(self, q,k, pad_ix_q, pad_ix_k):
        # x: [batch_size, seq_len]
        len_q, len_k = q.size(1), k.size(1)
        q = q.ne(pad_ix_q).unsqueeze(1).unsqueeze(3)
        q = q.repeat(1, 1, 1, len_k)

        k = k.ne(pad_ix_k).unsqueeze(1).unsqueeze(2)
        k = k.repeat(1, 1, len_q,1)
        mask = q & k
        return mask

    def make_casual_mask(self, q,k):
        len_q = q.size(1)
        len_k = k.size(1)
        mask = torch.tril(torch.ones(len_q, len_k)).type(torch.BoolTensor).to(self.device)
        return mask

    def forward(self,src, trg):
        src_mask = self.make_pad_mask(src, src, self.src_pad_ix, self.src_pad_ix)
        trg_mask = self.make_pad_mask(trg, trg, self.trg_pad_ix, self.trg_pad_ix) * self.make_casual_mask(trg, trg)
        # trg_mask = self.make_pad_mask(trg, trg)
        enc = self.encoder(src, src_mask)
        out = self.decoder(trg, src, trg_mask,src_mask)
        return out



def main():
    # 超参数
    src_pad_ix = 0
    trg_pad_ix = 0
    enc_voc_size = 100
    dec_voc_size = 100
    d_model = 32
    max_len = 20
    n_heads = 4
    ffn_hidden = 64
    n_layer = 2
    dropout = 0.1
    batch_size = 8
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    epochs = 10

    # 构造模型
    model = Transformer(
        src_pad_ix, trg_pad_ix, enc_voc_size, dec_voc_size, d_model,
        max_len, n_heads, ffn_hidden, n_layer, device, dropout
    ).to(device)
    criterion = torch.nn.CrossEntropyLoss(ignore_index=trg_pad_ix)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(epochs):
        model.train()
        # 生成模拟数据
        src = torch.randint(2, enc_voc_size, (batch_size, max_len)).to(device)
        trg = torch.randint(2, dec_voc_size, (batch_size, max_len)).to(device)
        # 目标输出右移一位
        trg_input = trg[:, :-1]
        trg_output = trg[:, 1:]

        # 前向传播
        output = model(src, trg_input)
        # output: [batch, seq_len, vocab_size]
        output = output.view(-1, dec_voc_size)
        trg_output = trg_output.reshape(-1)

        loss = criterion(output, trg_output)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}")

if __name__ == '__main__':
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    main()
